package com.lu.sb210408.gpu;

import com.sun.jna.Platform;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;

@Slf4j
public class Shell {

    private final GPU gpu = new GPU();

    public GPU getGpu() throws IOException {

        String gpuStr;

        if ((gpuStr = executeCommand()) == null) {

            return gpu;
        }

        // 1、获取Driver Version、CUDA Version
        String[] driverCUDA = gpuStr.split("\\+-------------------------------\\+----------------------\\+----------------------\\+")[0]
                .split("\\s{3,}");
        String driverVersion = driverCUDA[2].split(": ")[1];
        String CUDAVersion = driverCUDA[3].split(": ")[1];
        gpu.setDriverVersion(driverVersion);
        gpu.setCudaVersion(CUDAVersion);

        // 2、获取GPUInfo
        List<GPUInfo> GPUInfos = new ArrayList<>();
        String[] gpuInfos = gpuStr.split("\\|===============================\\+======================\\+======================\\|")[1]
                .split("                                                                               ")[0]
                .split("\\+-------------------------------\\+----------------------\\+----------------------\\+");
        for (int i = 0; i < gpuInfos.length - 1; i++) {
            GPUInfo gpuInfo = new GPUInfo();
            String[] firstRow = gpuInfos[i].split("\n")[1].split("\\s{2,}");
            gpuInfo.setNumber(firstRow[1]);
            gpuInfo.setName(firstRow[2]);
            String[] secondRow = gpuInfos[i].split("\n")[2].split("\\|");
            String[] firstColumn = secondRow[1].split("\\s{3,}");
            gpuInfo.setFan(firstColumn[0].replace(" ", ""));
            gpuInfo.setTemp(firstColumn[1].replace("C", "℃"));
            gpuInfo.setUsageCapPwr(firstColumn[3]);
            String[] secondColumn = secondRow[2].split("/");
            String GPUMemory = secondColumn[1].replace(" ", "");
            String usageGPUMemory = secondColumn[0].replace(" ", "");
            String unUsageGPUMemory = Integer.parseInt(GPUMemory.replace("MiB", "")) - Integer.parseInt(usageGPUMemory.replace("MiB", "")) + "MiB";
            double usageRate = Integer.parseInt(usageGPUMemory.replace("MiB", "")) * 100.00 / Integer.parseInt(GPUMemory.replace("MiB", ""));
            usageRate = ((int) (usageRate * 100)) * 1.0 / 100;
            gpuInfo.setGpuMemory(GPUMemory);
            gpuInfo.setUsageGpuMemory(usageGPUMemory);
            gpuInfo.setUnUsageGpuMemory(unUsageGPUMemory);
            gpuInfo.setUsageRate(usageRate);
            gpuInfo.setGpuRate(Double.parseDouble(secondRow[3].split("\\s{3,}")[1].replace("%", "")));
            GPUInfos.add(gpuInfo);
        }
        gpu.setGpuInfos(GPUInfos);

        // 3、获取Process
        List<Process> processes = new ArrayList<>();
        String[] processArray = gpuStr.split("\\|=============================================================================\\|")[1]
                .split("\n");
        for (int i = 0; i < processArray.length - 1; i++) {
            if (processArray[i].contains("No running processes found")) {
                break;
            }
            Process process = new Process();
            String[] pro = processArray[i].split("\\s+");
            if (pro.length < 6) {
                continue;
            }
            process.setNumber(pro[1]);
            process.setName(pro[4]);
            process.setUsageGPUMemory(pro[5]);
            processes.add(process);
        }
        gpu.setProcesses(processes);

        return gpu;
    }

    public String executeCommand() throws IOException {
        java.lang.Process process = null;
        try {
            if (Platform.isWindows()) {
                process = Runtime.getRuntime().exec("nvidia-smi.exe");
            } else if (Platform.isLinux()) {
                String[] shell = {"/bin/bash", "-c", "nvidia-smi"};
                process = Runtime.getRuntime().exec(shell);
            }
        } catch (IOException e) {
            log.warn("显卡不存在或获取显卡信息失败");
        }

        if (process == null) {
            return null;
        }

        process.getOutputStream().close();

        BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()));

        StringBuilder stringBuilder = new StringBuilder();

        String line;

        while (StringUtils.isNotEmpty(line = reader.readLine())) {

            stringBuilder.append(line).append("\n");
        }

        return stringBuilder.toString();
    }
}
